import torch
import torchvision.datasets as dsets
import numpy as np
from tqdm import tqdm

# MNIST dataset
train_dataset = dsets.MNIST(root='./pymnist',  # 选择数据的根目录
                            train=True,  # 选择训练集
                            transform=None,  # 不考虑使用任何数据预处理
                            download=True)  # 从网络上download图片
test_dataset = dsets.MNIST(root='./pymnist',  # 选择数据的根目录
                           train=False,  # 选择测试集
                           transform=None,  # 不考虑使用任何数据预处理
                           download=True)  # 从网络上download图片
labels = test_dataset.test_labels
finallabels = labels.reshape(labels.shape[0], 1)
batch_size = 2

print(finallabels[1:1 + batch_size])
one_hot = torch.zeros(batch_size, 10).scatter_(1, finallabels[1:1 + batch_size], 1)
print(one_hot)
